import logging
import functools
import json
from typing import List, Optional

from web3 import Web3
from web3.types import (
    Wei,
    ChecksumAddress,
)
from web3 import middleware
from web3.middleware import validation

from coin.exchange.uniswap_base.types import AddressLike
from coin.exchange.uniswap_base.token import (
    ERC20Token,
    Pool,
)
from coin.exchange.uniswap_base.exceptions import (
    InvalidToken,
    InvalidPool,
)
from coin.exchange.uniswap_base.util import (
    str_to_addr,
    addr_to_str,
    load_contract,
    load_contract_erc20,
    is_same_address,
)
from .constants import (
    UNISWAP_V2_FACTORY_CONTRACT_ADDRESS,
    UNISWAP_V2_ROUTER_CONTRACT_ADDRESS,
    UNISWAP_V3_QUOTER_CONTRACT_ADDRESS,
    UNISWAP_V3_ROUTER_CONTRACT_ADDRESS,
    ETH_ADDRESS,
    netid_to_name,
    SYMBOL_TO_ADDRESS,
    TOKEN_TO_ADDRESS,
    TOKEN_TO_DECIMAL,
)


def estimate_quote_qty(quote_address, base_address, blockchain_name):
  # Assume 5000 USD worth of quote. If quote is ETH, then worth of 2 ETH.
  if quote_address in (
      TOKEN_TO_ADDRESS[blockchain_name]['USDT'],
      TOKEN_TO_ADDRESS[blockchain_name]['DAI'],
      TOKEN_TO_ADDRESS[blockchain_name]['USDC']):
    if base_address in (TOKEN_TO_ADDRESS[blockchain_name]['CCX']):
      return 500
    else:
      return 5000
  elif quote_address == TOKEN_TO_ADDRESS[blockchain_name]['WETH']:
    if base_address in (TOKEN_TO_ADDRESS[blockchain_name]['MUSA']):
      return 0.25
    else:
      return 1
  elif quote_address == TOKEN_TO_ADDRESS[blockchain_name]['WBTC']:
    return 0.1
  else:
    raise NotImplementedError(quote_address)


class UniswapManager:
  def __init__(
      self,
      provider: str = None,
      web3: Web3 = None,
      version: int = 3,
      blockchain_name: str = 'ethereum',
      logger=None,
  ) -> None:
    """
    :param address: The public address of the ETH wallet to use.
    :param provider: Web3 provider URI.
    :param web3: Can be optionally set to a custom Web3 instance.
    :param version: Which version of the Uniswap contracts to use.
    :param factory_contract_addr:
    :param router_contract_addr: v2 only.
    """
    self._version = version
    self._blockchain_name = blockchain_name
    if web3 is not None:
      self._w3 = web3
    else:
      self._provider = provider
      self._w3 = Web3(Web3.HTTPProvider(self._provider, request_kwargs={"timeout": 60}))
      self._w3.middleware_onion.add(middleware.time_based_cache_middleware)
      self._w3.middleware_onion.add(middleware.latest_block_based_cache_middleware)
      self._w3.middleware_onion.add(middleware.simple_cache_middleware)
      validation.METHODS_TO_VALIDATE = []
    netid = int(self._w3.net.version)
    netname = netid_to_name[netid]
    self._logger = logger or logging.getLogger('UniswapManager')
    self._logger.info(f"Web3 client {self._w3} connected.")
    self._last_bid0_prices = {}
    self._last_ask0_prices = {}

    if self._version == 2:
      router_contract_addr = UNISWAP_V2_ROUTER_CONTRACT_ADDRESS[netname]
      self._router_contract_address: AddressLike = str_to_addr(router_contract_addr)

      factory_contract_addr = UNISWAP_V2_FACTORY_CONTRACT_ADDRESS[netname]
      self._factory_contract = load_contract(
          self._w3,
          abi_path="data/coin1/abi/uniswap_v2/factory.abi",
          address=str_to_addr(factory_contract_addr),
      )
      # Documented here: https://uniswap.org/docs/v2/smart-contracts/router02/
      self._router_contract = load_contract(
          self._w3,
          abi_path="data/coin1/abi/uniswap_v2/router02.abi",
          address=self._router_contract_address,
      )
    elif self._version == 3:
      router_contract_addr = UNISWAP_V3_ROUTER_CONTRACT_ADDRESS[netname]
      self._router_contract_address = str_to_addr(router_contract_addr)

      quoter_contract_addr = UNISWAP_V3_QUOTER_CONTRACT_ADDRESS[netname]
      self._quoter_contract = load_contract(
          self._w3,
          abi_path="data/coin1/abi/uniswap_v3/quoter.abi",
          address=quoter_contract_addr,
      )
      self._router_contract = load_contract(
          self._w3,
          abi_path="data/coin1/abi/uniswap_v3/router.abi",
          address=self._router_contract_address,
      )
    else:
      raise ValueError(f"Invalid version: {self._version}, only 2 or 3 supported!")
    self._pools = {}
    self._pool_event_filters = {}
    self._tx_hash_set = set()

  @functools.lru_cache()
  def get_weth_address(self) -> ChecksumAddress:
    """Retrieves the WETH address from the contracts (which may vary between chains)."""
    if self._version == 2:
      # Contract calls should always return checksummed addresses
      address: ChecksumAddress = self._router_contract.functions.WETH().call()
    elif self._version == 3:
      address = self._router_contract.functions.WETH9().call()
    else:
      raise ValueError(self._version)
    return address

  @functools.lru_cache()
  def _get_pool(self, symbol: str) -> Pool:
    """
    Retrieves metadata from the Pool contract,
    """
    if symbol not in SYMBOL_TO_ADDRESS[self._blockchain_name]:
      base, quote = symbol.split('-')
      base_address = TOKEN_TO_ADDRESS[self._blockchain_name][base]
      base_decimal = TOKEN_TO_DECIMAL[base]
      quote_address = TOKEN_TO_ADDRESS[self._blockchain_name][quote]
      quote_decimal = TOKEN_TO_DECIMAL[quote]
      token0 = ERC20Token(base, base_address, base, base_decimal)
      token1 = ERC20Token(quote, quote_address, quote, quote_decimal)
      return Pool(symbol, None, token0, token1, 3000, True)
    else:
      pool_contract_addr = SYMBOL_TO_ADDRESS[self._blockchain_name][symbol]
      if self._version == 2:
        abi_path = "data/coin1/abi/uniswap_v2/pair.abi"
      elif self._version == 3:
        abi_path = "data/coin1/abi/uniswap_v3/pool.abi"
      else:
        raise ValueError(self._version)
      pool_contract = load_contract(
          self._w3,
          abi_path=abi_path,
          address=pool_contract_addr,
      )
      try:
        token0 = self.get_token(pool_contract.functions.token0().call())
        token1 = self.get_token(pool_contract.functions.token1().call())
        if self._version == 2:
          fee = 3000
        elif self._version == 3:
          fee = pool_contract.functions.fee().call()
        else:
          raise ValueError(self._version)
        base, quote = symbol.split('-')
        if base in token0.symbol:
          base_then_quote = True
        else:
          base_then_quote = False
      except Exception as e:
        self._logger.error(
            f"Exception occurred while trying to get pool {symbol}: {e}")
        raise InvalidPool(pool_contract_addr)
      return Pool(symbol, pool_contract, token0, token1, fee, base_then_quote)

  @functools.lru_cache()
  def _get_pool_filter(self, symbol: str):
    pool = self._get_pool(symbol)
    self._pools[symbol] = pool
    if pool.contract is None:
      return None
    self._pool_event_filters[symbol] = self._pools[symbol].contract.events.Swap \
        .create_filter(fromBlock='latest')
    return self._pool_event_filters[symbol]

  def get_as_bid0(self, symbol: str, fee=500) -> list:
    pool = self._get_pool(symbol)
    if pool.base_then_quote:
      base_address = pool.token0.address
      base_decimals = pool.token0.decimals
      quote_address = pool.token1.address
      quote_decimals = pool.token1.decimals
    else:
      base_address = pool.token1.address
      base_decimals = pool.token1.decimals
      quote_address = pool.token0.address
      quote_decimals = pool.token0.decimals
    quote_qty = estimate_quote_qty(quote_address, base_address, self._blockchain_name)
    bid_qty = self.get_price_output(
        base_address, quote_address, int(quote_qty * 10**quote_decimals), fee=pool.fee)
    bid_qty /= 10**base_decimals
    bid_price = quote_qty / bid_qty
    self._last_bid0_prices[symbol] = bid_price
    return [bid_price, bid_qty]

  def get_as_ask0(self, symbol: str, fee=500) -> list:
    pool = self._get_pool(symbol)
    if pool.base_then_quote:
      base_address = pool.token0.address
      base_decimals = pool.token0.decimals
      quote_address = pool.token1.address
      quote_decimals = pool.token1.decimals
    else:
      base_address = pool.token1.address
      base_decimals = pool.token1.decimals
      quote_address = pool.token0.address
      quote_decimals = pool.token0.decimals
    quote_qty = estimate_quote_qty(quote_address, base_address, self._blockchain_name)
    ask_qty = self.get_price_input(
        quote_address, base_address, int(quote_qty * 10**quote_decimals), fee=pool.fee)
    ask_qty /= 10**base_decimals
    ask_price = quote_qty / ask_qty
    self._last_ask0_prices[symbol] = ask_price
    return [ask_price, ask_qty]

  def get_trades(self, symbol: str) -> list:
    pool = self._get_pool(symbol)
    filter = self._get_pool_filter(symbol)
    if filter is None:
      return []
    try:
      events = filter.get_new_entries()
    except ValueError as e:
      self._logger.error(e)
      # Try to reset filter.
      self._get_pool_filter.cache_clear()
      filter = self._get_pool_filter(symbol)
      events = filter.get_all_entries()
    event_jsons = []
    if len(events) >= 1:
      tx_hash = events[0]['transactionHash']
      if tx_hash not in self._tx_hash_set:
        self._tx_hash_set.add(tx_hash)
      else:
        return []

      last_bid0_price = self._last_bid0_prices.get(symbol, None)
      if last_bid0_price is not None:
        valid_price_floor = last_bid0_price * (1 - 0.2)
      else:
        # bbo price is broken, then the subscriber is damaged.
        # So just publish whatever fetched from trade if possible.
        valid_price_floor = 0

      last_ask0_price = self._last_ask0_prices.get(symbol, None)
      if last_ask0_price is not None:
        valid_price_ceil = last_ask0_price * (1 + 0.2)
      else:
        # bbo price is broken, then the subscriber is damaged.
        # So just publish whatever fetched from trade if possible.
        valid_price_ceil = float('inf')

      for swap_event in events:
        swap_event_json = json.loads(Web3.to_json(swap_event))
        if self._version == 2:
          amount0 = swap_event_json['args']['amount0Out']
          amount1 = swap_event_json['args']['amount1In']
          amount0_out = True
          if amount0 == 0 and amount1 == 0:
            amount0 = swap_event_json['args']['amount0In']
            amount1 = swap_event_json['args']['amount1Out']
            amount0_out = False
          pair = pool
          if pair.base_then_quote:
            base_amount = amount0 / (10 ** pair.token0.decimals)
            quote_amount = amount1 / (10 ** pair.token1.decimals)
            if amount0_out:
              base_out = True
            else:
              base_out = False
          else:
            base_amount = amount1 / (10 ** pair.token1.decimals)
            quote_amount = amount0 / (10 ** pair.token0.decimals)
            if amount0_out:
              base_out = False
            else:
              base_out = True
          if base_out:
            # Contract perspective is opposite to trader perspective.
            side = 'B'
          else:
            side = 'S'
        elif self._version == 3:
          amount0 = swap_event_json['args']['amount0']
          amount1 = swap_event_json['args']['amount1']
          if pool.base_then_quote:
            base_amount = amount0 / (10**pool.token0.decimals)
            quote_amount = amount1 / (10**pool.token1.decimals)
          else:
            base_amount = amount1 / (10**pool.token1.decimals)
            quote_amount = amount0 / (10**pool.token0.decimals)
          if base_amount > 0:
            # amount in and out is from the prespective of the pool.
            side = 'S'
            assert quote_amount < 0, quote_amount
            quote_amount = -quote_amount
          else:
            side = 'B'
            base_amount = -base_amount
        else:
          raise ValueError(self._version)
        price = quote_amount / base_amount
        if price > valid_price_floor and price < valid_price_ceil:
          event_jsons.append({
            'price': price,
            'qty': base_amount,
            'side': side,
          })
        else:
          self._logger.error(f"Filter out invalid price {symbol}: {price}")
    return event_jsons

  def get_price_input(
      self,
      token0: AddressLike,
      token1: AddressLike,
      qty: int,
      fee: int = None,
      route: Optional[List[AddressLike]] = None,
  ) -> int:
    """Returns the amount of the input token you get for `qty` of the output token"""
    if fee is None:
      fee = 3000  # 3000/1e6 => 0.3%

    if token0 == ETH_ADDRESS:
      return self._get_eth_token_input_price(token1, Wei(qty), fee)
    elif token1 == ETH_ADDRESS:
      return self._get_token_eth_input_price(token0, qty, fee)
    else:
      return self._get_token_token_input_price(token0, token1, qty, fee, route)

  def get_price_output(
      self,
      token0: AddressLike,
      token1: AddressLike,
      qty: int,
      fee: int = None,
      route: Optional[List[AddressLike]] = None,
  ) -> int:
    """Returns the amount of input token you need to get `qty` of the output token"""
    if fee is None:
      fee = 3000  # 3000/1e6 => 0.3%

    if is_same_address(token0, ETH_ADDRESS):
      return self._get_eth_token_output_price(token1, qty, fee)
    elif is_same_address(token1, ETH_ADDRESS):
      return self._get_token_eth_output_price(token0, Wei(qty), fee)
    else:
      return self._get_token_token_output_price(token0, token1, qty, fee, route)

  def _get_eth_token_input_price(self, token: AddressLike, qty: Wei, fee: int) -> Wei:
    """Public price for ETH to Token trades with an exact input."""
    if self._version == 2:
      price = self._router_contract.functions.getAmountsOut(
          qty, [self.get_weth_address(), token]).call()[-1]
    elif self._version == 3:
      price = self._get_token_token_input_price(
          self.get_weth_address(), token, qty, fee=fee)
    else:
      raise ValueError(self._version)
    return price

  def _get_token_eth_input_price(self, token: AddressLike, qty: int, fee: int) -> int:
    """Public price for token to ETH trades with an exact input."""
    if self._version == 2:
      price = self._router_contract.functions.getAmountsOut(
          qty, [token, self.get_weth_address()]).call()[-1]
    elif self._version == 3:
      price = self._get_token_token_input_price(
          token, self.get_weth_address(), qty, fee=fee)
    else:
      raise ValueError(self._version)
    return price

  def _get_token_token_input_price(
      self,
      token0: AddressLike,
      token1: AddressLike,
      qty: int,
      fee: int,
      route: Optional[List[AddressLike]] = None,
  ) -> int:
    """
    Public price for token to token trades with an exact input.
    :param fee: (v3 only) The pool's fee in hundredths of a bip, i.e. 1e-6 (3000 is 0.3%)
    """
    if route is None:
      if self._version == 2:
        # If one of the tokens are WETH, delegate to appropriate call.
        if is_same_address(token0, self.get_weth_address()):
          return int(self._get_eth_token_input_price(token1, Wei(qty), fee))
        elif is_same_address(token1, self.get_weth_address()):
          return int(self._get_token_eth_input_price(token0, qty, fee))

        route = [token0, self.get_weth_address(), token1]
        self._logger.debug(f"No route specified, assuming route: {route}")

    if self._version == 2:
      price: int = self._router_contract.functions.getAmountsOut(qty, route).call()[-1]
    elif self._version == 3:
      if route:
        # NOTE: to support custom routes we need to support the Path data encoding:
        # https://github.com/Uniswap/uniswap-v3-periphery/blob/main/contracts/libraries/Path.sol
        # result: tuple = self.quoter.functions.quoteExactInput(route, qty).call()
        raise NotImplementedError("Custom route not yet supported for v3!")

      # FIXME: How to calculate this properly?
      sqrtPriceLimitX96 = 0
      price = self._quoter_contract.functions.quoteExactInputSingle(
          token0, token1, fee, qty, sqrtPriceLimitX96).call()
    else:
      raise ValueError(self._version)
    return price

  def _get_eth_token_output_price(self, token: AddressLike, qty: int, fee: int = None) -> Wei:
    """Public price for ETH to Token trades with an exact output."""
    if self._version == 2:
      route = [self.get_weth_address(), token]
      price = self._router_contract.functions.getAmountsIn(qty, route).call()[0]
    elif self._version == 3:
      price = Wei(self._get_token_token_output_price(
          self.get_weth_address(), token, qty, fee=fee))
    else:
      raise ValueError(self._version)
    return price

  def _get_token_eth_output_price(self, token: AddressLike, qty: Wei, fee: int = None) -> int:
    """Public price for token to ETH trades with an exact output."""
    if self._version == 2:
      route = [token, self.get_weth_address()]
      price = self._router_contract.functions.getAmountsIn(qty, route).call()[0]
    elif self._version == 3:
      price = self._get_token_token_output_price(
          token, self.get_weth_address(), qty, fee=fee)
    else:
      raise ValueError(self._version)
    return price

  def _get_token_token_output_price(
      self,
      token0: AddressLike,
      token1: AddressLike,
      qty: int,
      fee: int = None,
      route: Optional[List[AddressLike]] = None,
  ) -> int:
    """
    Public price for token to token trades with an exact output.
    :param fee: (v3 only) The pool's fee in hundredths of a bip, i.e. 1e-6 (3000 is 0.3%)
    """
    if not route:
      if self._version == 2:
        # If one of the tokens are WETH, delegate to appropriate call.
        if is_same_address(token0, self.get_weth_address()):
          return int(self._get_eth_token_output_price(token1, qty, fee))
        elif is_same_address(token1, self.get_weth_address()):
          return int(self._get_token_eth_output_price(token0, Wei(qty), fee))

        route = [token0, self.get_weth_address(), token1]
        self._logger.debug(f"No route specified, assuming route: {route}")

    if self._version == 2:
      price: int = self._router_contract.functions.getAmountsIn(qty, route).call()[0]
    elif self._version == 3:
      if route:
        # NOTE: to support custom routes we need to support the Path data encoding:
        # https://github.com/Uniswap/uniswap-v3-periphery/blob/main/contracts/libraries/Path.sol
        # result: tuple = self.quoter.functions.quoteExactOutput(route, qty).call()
        raise NotImplementedError("Custom route not yet supported for v3!")

      # FIXME: How to calculate this properly?
      sqrtPriceLimitX96 = 0
      price = self._quoter_contract.functions.quoteExactOutputSingle(
          token0, token1, fee, qty, sqrtPriceLimitX96).call()
    else:
      raise ValueError(self._version)
    return price

  @functools.lru_cache()
  def get_token(self, address: AddressLike) -> ERC20Token:
    """
    Retrieves metadata from the ERC20 contract of a given token,
    like its name, symbol, and decimals.
    """
    token_contract = load_contract_erc20(
        self._w3,
        address=address)
    try:
      name = token_contract.functions.name().call()
      symbol = token_contract.functions.symbol().call()
      decimals = token_contract.functions.decimals().call()
    except Exception as e:
      self._logger.error(
          f"Exception occurred while trying to get token {addr_to_str(address)}: {e}")
      raise InvalidToken(address)
    return ERC20Token(symbol, address, name, decimals)
